{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2a54701",
   "metadata": {},
   "source": [
    "# 主要使用transformers库。 \n",
    "# 框架主要使用pytorch"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8adacf4b",
   "metadata": {},
   "source": [
    "# 处理数据 中文字符---》 数字\n",
    "# 创建数据集。 把处理好的数据变成pytorch的数据集。 \n",
    "# 生成模型， 有了transformers库， 一般不需要自己创建模型。\n",
    "# 训练预测过程。"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "fccd515b",
   "metadata": {},
   "source": [
    "### 处理数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "1919c521",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-25T18:34:51.920761Z",
     "start_time": "2023-12-25T18:34:51.905800Z"
    }
   },
   "source": [
    "# 配置代理\n",
    "import os\n",
    "\n",
    "os.environ['http_proxy'] = '127.0.0.1:10809'\n",
    "os.environ['https_proxy'] = '127.0.0.1:10809'"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "bb40b74e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-25T18:34:17.999022Z",
     "start_time": "2023-12-25T18:34:06.478834Z"
    }
   },
   "source": [
    "from transformers import AutoTokenizer"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "bc0ac7ae",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-25T18:34:55.966196Z",
     "start_time": "2023-12-25T18:34:54.560427Z"
    }
   },
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained('uer/gpt2-chinese-cluecorpussmall')"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c2c3d845",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-25T18:34:58.787107Z",
     "start_time": "2023-12-25T18:34:58.764169Z"
    },
    "scrolled": true
   },
   "source": [
    "tokenizer"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f53f6ca4",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-25T18:35:01.232891Z",
     "start_time": "2023-12-25T18:35:01.214939Z"
    },
    "scrolled": true
   },
   "source": [
    "# 编码试算\n",
    "tokenizer.batch_encode_plus([\n",
    "    '明朝驿使发,一夜絮征袍.素手抽针冷,那堪把剪刀.裁缝寄远道,几日到临洮.',\n",
    "    '长安一片月,万户捣衣声.秋风吹不尽,总是玉关情.何日平胡虏,良人罢远征.'\n",
    "])"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "999a1679",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-25T18:35:10.874103Z",
     "start_time": "2023-12-25T18:35:04.365512Z"
    }
   },
   "source": [
    "import torch\n",
    "\n",
    "\n",
    "class Dataset(torch.utils.data.Dataset):\n",
    "    def __init__(self):\n",
    "        with open('chinese_poems.txt', encoding='utf-8') as f:\n",
    "            lines = f.readlines()\n",
    "        lines = [i.strip() for i in lines]\n",
    "        \n",
    "        self.lines = lines\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.lines)\n",
    "    \n",
    "    def __getitem__(self, i):\n",
    "        return self.lines[i]\n",
    "    \n",
    "dataset = Dataset()\n",
    "len(dataset), dataset[0]"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "836c240b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-25T18:35:10.890061Z",
     "start_time": "2023-12-25T18:35:10.876098Z"
    }
   },
   "source": [
    "def collate_fn(batch):\n",
    "    # 把中文编码成数字\n",
    "    data = tokenizer.batch_encode_plus(batch, padding=True, \n",
    "                               truncation=True,\n",
    "                               max_length=512,\n",
    "                               return_tensors='pt')\n",
    "    data['labels'] = data['input_ids'].clone()\n",
    "    return data"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "fae51f8f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-25T18:35:10.905021Z",
     "start_time": "2023-12-25T18:35:10.892059Z"
    }
   },
   "source": [
    "# 数据加载器\n",
    "loader = torch.utils.data.DataLoader(\n",
    "    dataset=dataset,\n",
    "    batch_size=4,\n",
    "    collate_fn=collate_fn,\n",
    "    shuffle=True,\n",
    "    drop_last=True\n",
    ")"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "98fe3a72",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-25T18:35:12.156672Z",
     "start_time": "2023-12-25T18:35:12.130742Z"
    }
   },
   "source": [
    "for i, data in enumerate(loader):\n",
    "    break"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "ac0d6a3c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-23T19:24:02.687609Z",
     "start_time": "2023-12-23T19:24:02.669535Z"
    }
   },
   "source": [
    "i"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "316724dc",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-25T18:35:14.340832Z",
     "start_time": "2023-12-25T18:35:14.313903Z"
    },
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "source": [
    "data"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "8c9ce959",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-25T18:35:18.502699Z",
     "start_time": "2023-12-25T18:35:18.487739Z"
    }
   },
   "source": [
    "len(data['input_ids'])"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "0c4a37cf",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-25T18:35:18.675239Z",
     "start_time": "2023-12-25T18:35:18.658283Z"
    },
    "scrolled": true
   },
   "source": [
    "for k, v in data.items():\n",
    "    print(k, v.shape)"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "7be8ebb7",
   "metadata": {},
   "source": [
    "### 创建模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "3d15b0fe",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-25T18:35:20.987010Z",
     "start_time": "2023-12-25T18:35:20.917196Z"
    }
   },
   "source": [
    "from transformers import AutoModelForCausalLM, GPT2Model"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "20f379be",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-25T18:35:25.249423Z",
     "start_time": "2023-12-25T18:35:23.047498Z"
    }
   },
   "source": [
    "# 加载模型\n",
    "model = AutoModelForCausalLM.from_pretrained('uer/gpt2-chinese-cluecorpussmall')\n"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "896e9f43",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-25T18:35:29.467142Z",
     "start_time": "2023-12-25T18:35:29.456171Z"
    }
   },
   "source": [
    "# 参数量\n",
    "sum(i.numel() for i in model.parameters())"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "11324a83",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-25T18:35:32.939853Z",
     "start_time": "2023-12-25T18:35:32.006350Z"
    }
   },
   "source": [
    "with torch.no_grad():\n",
    "    out = model(**data)"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "82092e53",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-25T18:35:34.893627Z",
     "start_time": "2023-12-25T18:35:34.882657Z"
    }
   },
   "source": [
    "out['loss'], out['logits'].shape"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "1becf6d9",
   "metadata": {},
   "source": [
    "### 训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "2350b277",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-25T18:35:37.002987Z",
     "start_time": "2023-12-25T18:35:36.995008Z"
    }
   },
   "source": [
    "from transformers import AdamW\n",
    "from transformers.optimization import get_scheduler"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "75b0f36c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-25T18:35:44.087202Z",
     "start_time": "2023-12-25T18:35:44.068253Z"
    }
   },
   "source": [
    "def train():\n",
    "    global model\n",
    "    device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "    model = model.to(device)\n",
    "    \n",
    "    optimizer = AdamW(model.parameters(), lr=5e-5)\n",
    "    scheduler = get_scheduler(name='linear',\n",
    "                             num_warmup_steps=0,\n",
    "                             num_training_steps=len(loader),\n",
    "                             optimizer=optimizer)\n",
    "    \n",
    "    model.train()\n",
    "    for i, data in enumerate(loader):\n",
    "        for k in data.keys():\n",
    "            data[k] = data[k].to(device)\n",
    "            \n",
    "        out = model(**data)\n",
    "        loss = out['loss']\n",
    "        \n",
    "        loss.backward()\n",
    "        # 梯度裁剪\n",
    "        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
    "        \n",
    "        optimizer.step()\n",
    "        scheduler.step()\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "        model.zero_grad()\n",
    "        \n",
    "        if i % 1000 == 0:\n",
    "            labels = data['labels'][:, 1:]\n",
    "            out = out['logits'].argmax(dim=2)[:, :-1]\n",
    "            select = labels != 0\n",
    "            labels = labels[select]\n",
    "            out = out[select]\n",
    "            del select \n",
    "            \n",
    "            # 计算准确率\n",
    "            accuracy = (labels == out).sum().item() / labels.numel()\n",
    "            lr = optimizer.state_dict()['param_groups'][0]['lr']\n",
    "            print(i, loss.item(), lr, accuracy)"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "15690d58",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-25T18:42:05.627065Z",
     "start_time": "2023-12-25T18:35:46.765250Z"
    },
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "source": [
    "train()"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b01c8e1f",
   "metadata": {},
   "source": [
    "# 保存模型\n",
    "# model = model.to('cpu')\n",
    "# torch.save(model, 'model.pt')"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "a6d68e13",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-25T19:05:05.520483Z",
     "start_time": "2023-12-25T19:05:05.506521Z"
    }
   },
   "source": [
    "def generate(text, row, col, model):\n",
    "    def generate_loop(data):\n",
    "        with torch.no_grad():\n",
    "            out = model(**data)\n",
    "        \n",
    "        # [5, b, vocab_size]\n",
    "        out = out['logits']\n",
    "        # [5, vocab_size]\n",
    "        out = out[:, -1]\n",
    "        \n",
    "        # [5, vocab_size] -> [5, 50]\n",
    "        topk_value = torch.topk(out, 50).values\n",
    "        # [5, 50], [5] -> [5, 1]\n",
    "        topk_value = topk_value[:, -1].unsqueeze(dim=1)\n",
    "        \n",
    "        # 赋值\n",
    "        out = out.masked_fill(out < topk_value, -float('inf'))\n",
    "        \n",
    "        # 不允许写特殊字符\n",
    "        out[:, tokenizer.sep_token_id] = -float('inf')\n",
    "        out[:, tokenizer.unk_token_id] = -float('inf')\n",
    "        out[:, tokenizer.pad_token_id] = -float('inf')\n",
    "        \n",
    "        for i in '，。':\n",
    "            out[:, tokenizer.get_vocab()[i]] = -float('inf')\n",
    "            \n",
    "        # [5, vocab_size] -> [5, 1]\n",
    "        out = out.softmax(dim=1)\n",
    "        out = out.multinomial(num_samples=1)\n",
    "        \n",
    "        # 强制添加标点符号\n",
    "        c = data['input_ids'].shape[1] / (col + 1)\n",
    "        if c % 1 == 0:\n",
    "            if c % 2 == 0:\n",
    "                out[:, 0] = tokenizer.get_vocab()['。']\n",
    "            else:\n",
    "                out[:, 0] = tokenizer.get_vocab()['，']\n",
    "        \n",
    "        data['input_ids'] = torch.cat([data['input_ids'], out], dim=1)\n",
    "        data['attention_mask'] = torch.ones_like(data['input_ids'])\n",
    "        data['token_type_ids'] = torch.zeros_like(data['input_ids'])\n",
    "        data['labels'] = data['input_ids'].clone()\n",
    "        \n",
    "        if data['input_ids'].shape[1] >= row * col + row + 1:\n",
    "            return data\n",
    "        return generate_loop(data)\n",
    "    \n",
    "    # 重复三遍\n",
    "    data = tokenizer.batch_encode_plus([text] * 3, return_tensors='pt')\n",
    "    data['input_ids'] = data['input_ids'][:, :-1]\n",
    "    data['attention_mask'] = torch.ones_like(data['input_ids'])\n",
    "    data['token_type_ids'] = torch.zeros_like(data['input_ids'])\n",
    "    data['labels'] = data['input_ids'].clone()\n",
    "    \n",
    "    data = generate_loop(data)\n",
    "    \n",
    "    for i in range(3):\n",
    "        print(i, tokenizer.decode(data['input_ids'][i]))"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "ce0c15f6",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-25T19:04:44.356816Z",
     "start_time": "2023-12-25T19:04:44.138400Z"
    }
   },
   "source": "model = torch.load('al_poetry.model')",
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "61afaf5c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-25T19:05:46.697299Z",
     "start_time": "2023-12-25T19:05:43.733228Z"
    }
   },
   "source": [
    "generate('秋高气爽', row=4, col=7, model=model)"
   ],
   "outputs": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.4"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
